#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
混合检索模块：关键词检索 + 向量检索
"""

import os
import json
import chromadb
from typing import List, Dict, Any, Tuple
from elasticsearch import Elasticsearch
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
from rag.emm.text_rag import TextEmbedding
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class HybridRetrieval:
    """混合检索器"""
    
    def __init__(self, 
                 collection_name: str = "nature_human_relationship",
                 es_index_name: str = "nature_documents",
                 es_host: str = "http://localhost:9200"):
        """
        初始化混合检索器
        
        Args:
            collection_name: ChromaDB集合名称
            es_index_name: Elasticsearch索引名称
            es_host: Elasticsearch主机地址
        """
        self.collection_name = collection_name
        self.es_index_name = es_index_name
        self.es_host = es_host
        
        # 初始化组件
        self.embedding_model = TextEmbedding()
        
        # 初始化ChromaDB
        self.chroma_client = chromadb.PersistentClient(path="./chromadb_nature")
        self.collection = None
        
        # 初始化Elasticsearch
        self.es_client = None
        
        # 权重配置
        self.vector_weight = 0.7  # 向量检索权重
        self.keyword_weight = 0.3  # 关键词检索权重
        
    def init_elasticsearch(self) -> bool:
        """
        初始化Elasticsearch连接和索引
        
        Returns:
            bool: 初始化是否成功
        """
        try:
            self.es_client = Elasticsearch([self.es_host])
            
            # 检查连接
            if not self.es_client.ping():
                logger.warning("Elasticsearch连接失败，将只使用向量检索")
                return False
            
            # 创建索引（如果不存在）
            if not self.es_client.indices.exists(index=self.es_index_name):
                mapping = {
                    "mappings": {
                        "properties": {
                            "id": {"type": "keyword"},
                            "text": {"type": "text", "analyzer": "ik_max_word"},
                            "chunk_index": {"type": "integer"},
                            "chunk_size": {"type": "integer"},
                            "source": {"type": "keyword"},
                            "created_at": {"type": "date"}
                        }
                    }
                }
                self.es_client.indices.create(index=self.es_index_name, body=mapping)
                logger.info(f"创建Elasticsearch索引: {self.es_index_name}")
            
            return True
            
        except Exception as e:
            logger.error(f"Elasticsearch初始化失败: {e}")
            return False
    
    def init_chromadb(self) -> bool:
        """
        初始化ChromaDB连接
        
        Returns:
            bool: 初始化是否成功
        """
        try:
            self.collection = self.chroma_client.get_collection(self.collection_name)
            logger.info(f"成功连接ChromaDB集合: {self.collection_name}")
            return True
        except Exception as e:
            logger.error(f"ChromaDB初始化失败: {e}")
            return False
    
    def vector_search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        向量检索
        
        Args:
            query: 查询文本
            top_k: 返回结果数量
            
        Returns:
            List[Dict]: 检索结果列表
        """
        try:
            if not self.collection:
                if not self.init_chromadb():
                    return []
            
            # 生成查询向量
            embeddings = self.embedding_model.generate_embeddings(query)
            query_embedding = embeddings[0] if embeddings else []
            
            # 确保embedding是列表格式
            if isinstance(query_embedding, list) and len(query_embedding) > 0:
                if isinstance(query_embedding[0], list):
                    query_embedding = query_embedding[0]
            
            # 执行向量检索
            results = self.collection.query(
                query_embeddings=[query_embedding],
                n_results=top_k
            )
            
            # 格式化结果
            vector_results = []
            if results['ids'] and results['ids'][0]:
                for i in range(len(results['ids'][0])):
                    result = {
                        "id": results['ids'][0][i],
                        "text": results['documents'][0][i],
                        "score": 1 - results['distances'][0][i],  # 转换为相似度分数
                        "metadata": results['metadatas'][0][i],
                        "retrieval_type": "vector"
                    }
                    vector_results.append(result)
            
            logger.info(f"向量检索完成，返回 {len(vector_results)} 个结果")
            return vector_results
            
        except Exception as e:
            logger.error(f"向量检索失败: {e}")
            return []
    
    def keyword_search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        关键词检索
        
        Args:
            query: 查询文本
            top_k: 返回结果数量
            
        Returns:
            List[Dict]: 检索结果列表
        """
        try:
            if not self.es_client:
                if not self.init_elasticsearch():
                    return []
            
            # 构建查询
            search_body = {
                "query": {
                    "multi_match": {
                        "query": query,
                        "fields": ["text^2", "source"],
                        "type": "best_fields",
                        "fuzziness": "AUTO"
                    }
                },
                "size": top_k,
                "_source": ["id", "text", "chunk_index", "chunk_size", "source", "created_at"]
            }
            
            # 执行搜索
            response = self.es_client.search(
                index=self.es_index_name,
                body=search_body
            )
            
            # 格式化结果
            keyword_results = []
            for hit in response['hits']['hits']:
                result = {
                    "id": hit['_source']['id'],
                    "text": hit['_source']['text'],
                    "score": hit['_score'] / 10.0,  # 归一化分数
                    "metadata": {
                        "chunk_index": hit['_source'].get('chunk_index'),
                        "chunk_size": hit['_source'].get('chunk_size'),
                        "source": hit['_source'].get('source'),
                        "created_at": hit['_source'].get('created_at')
                    },
                    "retrieval_type": "keyword"
                }
                keyword_results.append(result)
            
            logger.info(f"关键词检索完成，返回 {len(keyword_results)} 个结果")
            return keyword_results
            
        except Exception as e:
            logger.error(f"关键词检索失败: {e}")
            return []
    
    def hybrid_search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        混合检索：结合向量检索和关键词检索
        
        Args:
            query: 查询文本
            top_k: 返回结果数量
            
        Returns:
            List[Dict]: 混合检索结果列表
        """
        logger.info(f"开始混合检索，查询: {query}")
        
        # 执行向量检索
        vector_results = self.vector_search(query, top_k * 2)
        
        # 执行关键词检索
        keyword_results = self.keyword_search(query, top_k * 2)
        
        # 合并和重排序结果
        combined_results = self._combine_results(vector_results, keyword_results, top_k)
        
        logger.info(f"混合检索完成，返回 {len(combined_results)} 个结果")
        return combined_results
    
    def _combine_results(self, 
                        vector_results: List[Dict[str, Any]], 
                        keyword_results: List[Dict[str, Any]], 
                        top_k: int) -> List[Dict[str, Any]]:
        """
        合并和重排序检索结果
        
        Args:
            vector_results: 向量检索结果
            keyword_results: 关键词检索结果
            top_k: 返回结果数量
            
        Returns:
            List[Dict]: 合并后的结果列表
        """
        # 创建结果字典，用于去重和合并
        result_dict = {}
        
        # 添加向量检索结果
        for result in vector_results:
            doc_id = result['id']
            if doc_id not in result_dict:
                result_dict[doc_id] = {
                    "id": doc_id,
                    "text": result['text'],
                    "metadata": result['metadata'],
                    "vector_score": result['score'] * self.vector_weight,
                    "keyword_score": 0.0,
                    "combined_score": result['score'] * self.vector_weight,
                    "retrieval_types": ["vector"]
                }
            else:
                result_dict[doc_id]["vector_score"] = result['score'] * self.vector_weight
                result_dict[doc_id]["retrieval_types"].append("vector")
        
        # 添加关键词检索结果
        for result in keyword_results:
            doc_id = result['id']
            if doc_id not in result_dict:
                result_dict[doc_id] = {
                    "id": doc_id,
                    "text": result['text'],
                    "metadata": result['metadata'],
                    "vector_score": 0.0,
                    "keyword_score": result['score'] * self.keyword_weight,
                    "combined_score": result['score'] * self.keyword_weight,
                    "retrieval_types": ["keyword"]
                }
            else:
                result_dict[doc_id]["keyword_score"] = result['score'] * self.keyword_weight
                result_dict[doc_id]["retrieval_types"].append("keyword")
        
        # 计算综合分数
        for doc_id, result in result_dict.items():
            result["combined_score"] = result["vector_score"] + result["keyword_score"]
        
        # 按综合分数排序
        sorted_results = sorted(result_dict.values(), 
                              key=lambda x: x["combined_score"], 
                              reverse=True)
        
        # 格式化最终结果
        final_results = []
        for result in sorted_results[:top_k]:
            final_result = {
                "id": result["id"],
                "text": result["text"],
                "score": result["combined_score"],
                "metadata": result["metadata"],
                "retrieval_types": result["retrieval_types"],
                "vector_score": result["vector_score"],
                "keyword_score": result["keyword_score"]
            }
            final_results.append(final_result)
        
        return final_results
    
    def index_documents_to_elasticsearch(self, chunk_docs: List[Dict[str, Any]]) -> bool:
        """
        将文档索引到Elasticsearch
        
        Args:
            chunk_docs: 文档块列表
            
        Returns:
            bool: 索引是否成功
        """
        try:
            if not self.es_client:
                if not self.init_elasticsearch():
                    return False
            
            # 批量索引文档
            from elasticsearch.helpers import bulk
            
            actions = []
            for doc in chunk_docs:
                action = {
                    "_index": self.es_index_name,
                    "_id": doc["id"],
                    "_source": {
                        "id": doc["id"],
                        "text": doc["text"],
                        "chunk_index": doc["chunk_index"],
                        "chunk_size": doc["chunk_size"],
                        "source": doc["source"],
                        "created_at": doc["metadata"].get("created_at", "2024-01-01")
                    }
                }
                actions.append(action)
            
            # 执行批量索引
            success, failed = bulk(self.es_client, actions)
            
            logger.info(f"Elasticsearch索引完成，成功: {success}, 失败: {failed}")
            return failed == 0
            
        except Exception as e:
            logger.error(f"Elasticsearch索引失败: {e}")
            return False


def main():
    """测试混合检索功能"""
    retrieval = HybridRetrieval()
    
    # 测试查询
    test_queries = [
        "人与自然的关系",
        "生态危机",
        "可持续发展",
        "工业革命的影响"
    ]
    
    for query in test_queries:
        print(f"\n查询: {query}")
        print("-" * 50)
        
        results = retrieval.hybrid_search(query, top_k=3)
        
        for i, result in enumerate(results, 1):
            print(f"{i}. 分数: {result['score']:.3f}")
            print(f"   类型: {result['retrieval_types']}")
            print(f"   文本: {result['text'][:100]}...")
            print()


if __name__ == "__main__":
    main()
